import platgo as pg
import numpy as np


class Rosebrock(pg.Problem):

    def __init__(self, D: int = 20) -> None:
        self.name = 'Rosenbrock'
        self.type['single'], self.type['real'], self.type['large'], self.type['expensive'] = [True] * 4
        self.M = 1
        self.D = D
        lb = [-5] * self.D
        ub = [10] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        x = pop.decs
        x1 = x[:, 0:len(x[0])-1]
        x2 = x[:, 1:len(x[0])]
        pop.objv = np.array([np.sum((100 * np.square(x1 - np.square(x2))), 1) + np.sum(np.square(x2 - 1), 1)]).T

    def get_optimal(self) -> np.ndarray:
        return np.array([[0]])